#include <stdio.h>
#include <stdlib.h>
#include <pcap.h>
#include <signal.h>
#include <inttypes.h>
#include <net/ethernet.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <net/ethernet.h>
#include <string.h>

void pcap_callback(u_char *user, const struct pcap_pkthdr *h, const u_char *bytes);

typedef enum { UDP_STAT, TCP_STAT, OTHER_STAT } stat_type;


typedef struct {
    uint64_t n_bytes;
    uint64_t n_packets;
} bandwidth_stat;

typedef struct {
    stat_type type;
    uint16_t dest_port;
    bandwidth_stat stat;
} pcap_stat_unit;

typedef struct {
    pcap_stat_unit *stat_units;
    unsigned n_stat_units, stat_units_allocated;
    bandwidth_stat slice_total;
    struct timeval timestamp;
} pcap_stat_slice;


typedef struct {
    struct timeval last_ts;
    pcap_stat_slice *slices;
    unsigned n_slices, slices_allocated;

    bandwidth_stat total, short_eth_packets, non_ip_packets, short_ip_packets;
} pcap_collected_data;

pcap_collected_data pcap_data;

void pcap_stop(int signum);

struct sigaction stop_pcap_action = {
    .sa_handler = pcap_stop
};

pcap_t *pcap_handle;

void init_pcap_data(pcap_collected_data *pcap_data);
void adjust_current_slice(pcap_collected_data *pcap_data, const struct timeval *ts);
void account_ip(pcap_collected_data *data, const u_char *bytes, uint32_t len, 
                uint32_t total_eth_len);
void account_ip6(pcap_collected_data *data, const u_char *bytes, uint32_t len,
                 uint32_t total_eth_len);

void account_packet(pcap_collected_data *data, stat_type type, uint16_t port, unsigned length);

void print_slice(pcap_stat_slice *slice);
        
int main(int argc, char **argv)
{
    char *dev, errbuf[PCAP_ERRBUF_SIZE];
    

    sigaction(SIGINT, &stop_pcap_action, NULL);

    init_pcap_data(&pcap_data);
    if (argc == 1) {
        dev = pcap_lookupdev(errbuf);
        if (dev == NULL) {
            fprintf(stderr, "Couldn't find default device: %s\n", errbuf);
            return(2);
        }
        printf("Default device: %s\n", dev);
    } else {
        dev = argv[1];
    }

    pcap_handle = pcap_open_live(dev, BUFSIZ, 0, 0, errbuf);
    if (!pcap_handle) {
        fprintf(stderr, "Couldn't open device %s: %s\n", dev, errbuf);
        return 2;
    }

    if (pcap_datalink(pcap_handle) != DLT_EN10MB) {
        fprintf(stderr, "Device %s doesn't provide Ethernet headers - not supported\n", dev);
        return(2);
    }

    pcap_loop(pcap_handle, -1, pcap_callback, (u_char *)&pcap_data);

    printf("Total: %" PRIu64 " packets, %" PRIu64 " bytes\n", pcap_data.total.n_packets, pcap_data.total.n_bytes);
    printf("Non-IP: %" PRIu64 " packets, %" PRIu64 " bytes\n", pcap_data.non_ip_packets.n_packets, pcap_data.non_ip_packets.n_bytes);
    printf("Discared Ethernet short frames: %" PRIu64 " packets, %" PRIu64 " bytes\n", pcap_data.short_eth_packets.n_packets,
           pcap_data.short_eth_packets.n_bytes);

    return 0;
}

void pcap_callback(u_char *user, const struct pcap_pkthdr *h, const u_char *bytes)
{
    pcap_collected_data *data = (pcap_collected_data *)user;
    const u_char *p = bytes;
    struct ether_header *eth_header;

    adjust_current_slice(data, &h->ts);
    
    data->total.n_packets++;
    data->total.n_bytes += h->len;

    data->slices[data->n_slices].slice_total.n_packets++;
    data->slices[data->n_slices].slice_total.n_bytes += h->len;

    eth_header = (struct ether_header *)bytes;
    p += sizeof(*eth_header);
    if ((unsigned)(p - bytes) > h->caplen) {
        data->short_eth_packets.n_packets++;
        data->short_eth_packets.n_bytes += h->len;
        return;
    }


    switch (ntohs(eth_header->ether_type)) {
    case ETHERTYPE_IP:
        account_ip(data, p, h->caplen - sizeof(*eth_header), h->len);
        break;
    case ETHERTYPE_IPV6:
        account_ip6(data, p, h->caplen - sizeof(*eth_header), h->len);
        break;
    default:
        data->non_ip_packets.n_packets++;
        data->non_ip_packets.n_bytes += h->len;
        break;
    }
}

void pcap_stop(int signum)
{
    pcap_breakloop(pcap_handle);
}

void account_ip(pcap_collected_data *data, const u_char *bytes, uint32_t len,
                uint32_t total_eth_len)
{
    struct ip *ip_header;
    const u_char *p;
    uint16_t dest_port;

    if (len < 20) {
        data->short_ip_packets.n_packets++;
        data->short_ip_packets.n_bytes += total_eth_len;
        return;
    }
        
    ip_header = (struct ip *)bytes;

    if (len < ntohs(ip_header->ip_len)) {
        data->short_ip_packets.n_packets++;
        data->short_ip_packets.n_bytes += len;
        return;
    }
    
    p = bytes + ip_header->ip_hl * 4;

    switch (ip_header->ip_p) {
    case IPPROTO_UDP:
        dest_port = htons(*((uint16_t *)(p + 2)));
        account_packet(data, UDP_STAT , dest_port, total_eth_len);
    case IPPROTO_TCP:
        dest_port = htons(*((uint16_t *)(p + 2)));
        account_packet(data, TCP_STAT , dest_port, total_eth_len);
        break;
    default:
        account_packet(data, OTHER_STAT, 0, total_eth_len);
        break;
    }
}

void account_ip6(pcap_collected_data *data, const u_char *bytes, uint32_t len,
                 uint32_t total_eth_len)
{
}

void init_pcap_data(pcap_collected_data *pcap_data)
{
    pcap_data->slices_allocated = 16;
    pcap_data->n_slices = 0;
    pcap_data->slices = calloc(pcap_data->slices_allocated, sizeof(*pcap_data->slices));
    gettimeofday(&pcap_data->last_ts, 0);
}

void adjust_current_slice(pcap_collected_data *pcap_data, const struct timeval *ts)
{
    unsigned time_diff;
    int ms_diff = ts->tv_usec - pcap_data->last_ts.tv_usec;
    pcap_stat_slice *old_slice, *new_slice;
    int i;

    if (ms_diff < 0) {
        ms_diff = 1000000 - ms_diff;
    }

    time_diff = (ts->tv_sec - pcap_data->last_ts.tv_sec) * 1000000 + ms_diff;
    time_diff /= 1000000;

    if (time_diff >   pcap_data->n_slices) {
        printf("Second passed\n");
        print_slice(pcap_data->slices + pcap_data->n_slices);
        if (time_diff >= pcap_data->slices_allocated) {
            pcap_data->slices_allocated *= 2;
            pcap_data->slices = realloc(pcap_data->slices, pcap_data->slices_allocated * sizeof(*pcap_data->slices));
            memset(pcap_data->slices + pcap_data->n_slices, 0, (pcap_data->slices_allocated - pcap_data->n_slices) * sizeof(*pcap_data->slices));
            printf("Reallocated to %u\n", pcap_data->slices_allocated);

        }
        pcap_data->n_slices = time_diff;
    }
}

const char *stat_type_to_str(stat_type type) 
{
    switch (type) {
    case UDP_STAT:
        return "UDP";
        break;
    case TCP_STAT:
        return "TCP";
        break;
    case OTHER_STAT:
        return "Other";
        break;
    }
    return "Unexpected";
}

void account_packet(pcap_collected_data *data, stat_type type, uint16_t port, unsigned length)
{
    //printf("%s packet to port %u, length %u\n", stat_type_to_str(type), port, length);
}

void print_slice(pcap_stat_slice *slice)
{
    printf("%" PRIu64 " packets, %" PRIu64 " bytes\n", slice->slice_total.n_packets, slice->slice_total.n_bytes);
}
